import tushare as ts
import backtrader as bt
import pandas as pd

# 1. 获取数据（以贵州茅台为例）
ts.set_token('b91f8e85973449023ab471c313af892fed91df1fab7fc669abc3de8f')  # 替换为实际Token[2](@ref)
pro = ts.pro_api()
df = pro.daily(ts_code='000063.SZ', start_date='20200101', end_date='20250324')

# 2. 数据格式转换（适配Backtrader要求）
df['trade_date'] = pd.to_datetime(df['trade_date'])
df = df.rename(columns={
    'trade_date': 'date',
    'open': 'open',
    'high': 'high',
    'low': 'low',
    'close': 'close',
    'vol': 'volume'
}).set_index('date').sort_index()

# 3. 创建Backtrader策略类
class MaCrossStrategy(bt.Strategy):
    params = (
        ('short_period', 5),   # 短期均线周期
        ('long_period', 20),  # 长期均线周期
    )

    def __init__(self):
        # 计算均线指标
        self.short_ma = bt.indicators.SMA(self.data.close, period=self.p.short_period)
        self.long_ma = bt.indicators.SMA(self.data.close, period=self.p.long_period)
        self.crossover = bt.indicators.CrossOver(self.short_ma, self.long_ma)  # 金叉死叉信号

    def next(self):
        if not self.position:  # 无持仓时
            if self.crossover > 0:  # 金叉信号
                self.order = self.buy(size=100)  # 买入100股
        elif self.crossover < 0:     # 死叉信号
            self.order = self.sell(size=100)   # 卖出全部持仓

# 4. 初始化回测引擎
cerebro = bt.Cerebro()
data = bt.feeds.PandasData(dataname=df)
cerebro.adddata(data)
cerebro.addstrategy(MaCrossStrategy)
cerebro.broker.setcash(100000.0)  # 初始资金10万元
cerebro.broker.setcommission(commission=0.001)  # 佣金率0.1%

# 5. 运行回测并输出结果
print('初始资金: %.2f' % cerebro.broker.getvalue())
cerebro.run()
print('最终资金: %.2f' % cerebro.broker.getvalue())
cerebro.plot()